from langchain import hub
from langchain.agents import AgentExecutor, create_react_agent,create_structured_chat_agent
import os
from dotenv import load_dotenv
from langchain_community.llms import Tongyi

os.environ["DASHSCOPE_API_KEY"] = "sk-9d8f1914800e497f8717144e860f99bc"

model = Tongyi(temperature=1)
from langchain.pydantic_v1 import BaseModel, Field
from langchain.tools import BaseTool, StructuredTool, tool
from langchain.agents import AgentExecutor, create_structured_chat_agent


class CalculatorInput(BaseModel):
    s: str = Field(description="输入字符串")


def multiply(s: str) -> int:
    """Multiply two numbers."""
    return len(s)


calculator = StructuredTool.from_function(
    func=multiply,  # 工具具体逻辑
    name="Calculator",  # 工具名
    description="计算字符长度",  # 工具信息
    args_schema=CalculatorInput,  # 工具接受参数信息
    return_direct=True,  # 直接作为工具的输出返回给调用者
    handle_tool_error=True,  # 报错了继续执行，不会吧那些报错行抛出，也可以自定义函数处理，handle_tool_error=函数名
)


class SearchInput(BaseModel):
    query: str = Field(description="should be a search query")


@tool("search-tool", args_schema=SearchInput, return_direct=True)
def search(query: str) -> str:
    """Look up things online."""
    return "你好啊"


tools = [search, calculator]

# 修改提示词
from langchain.prompts import PromptTemplate

prompt_template = """
Answer the following questions as best you can. You have access to the following tools:
{tools}
Use the following format:
Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question
Begin!
Question: {input}
Thought:{agent_scratchpad}
"""
prompt = PromptTemplate.from_template(prompt_template)

agent = create_react_agent(model, tools, prompt)

# 查看提示词
prompt_template = agent.get_prompts()[0]
print(prompt_template.format(input="what's 4.1*7.9=?", agent_scratchpad=""))

agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, return_intermediate_steps=True)
agent_executor.invoke({"input": "langchian是什么东西？`阿萨德防守打法`有多少个字符？"})

# import asyncio
#
#
# async def run():
#     response = await agent_executor.ainvoke({"input": "`阿萨德防守打法`有多少个字符？langchian是什么东西？"})
#     print(response)


# if __name__ == '__main__':
#     asyncio.run(run())

